"""
K-means 算法具体实现
"""
import numpy as np


class KMeans:
    def __init__(self, data, num_clusters):
        self.data = data
        self.num_clusters = num_clusters

    def train(self, max_iterations):
        # TODO 1. 初始化 K 个质心
        centroids = KMeans.centroids_init(self.data, self.num_clusters)
        # TODO 2. 开始训练
        num_examples = self.data.shape[0]
        # np.empty() 快速创建未初始化的数组
        # 相比 np.zeros() 或 np.ones()，np.empty() 跳过初始化步骤，直接分配内存，因此在处理大型数组时速度更快。
        closest_centroids_ids = np.empty((num_examples, 1))
        for _ in range(max_iterations):
            # TODO 3. 得到当前每一个样本点到 K 个质心的距离，找出最近的
            closest_centroids_ids = KMeans.centroids_find_closest(self.data, centroids)
            # TODO 4. 进行质心更新
            centroids = KMeans.centroids_compute(self.data, closest_centroids_ids, self.num_clusters)
        return centroids, closest_centroids_ids

    @staticmethod
    def centroids_init(data, num_clusters):
        """ 初始化（随机） num_clusters 个质心 """
        num_examples = data.shape[0]
        random_ids = np.random.permutation(num_examples)
        centroids = data[random_ids[:num_clusters], :]
        return centroids

    @staticmethod
    def centroids_find_closest(data, centroids):
        """ 计算每个样本点到质心的距离并找出最近的 """
        num_examples = data.shape[0]
        num_centroids = centroids.shape[0]
        closest_centroids_ids = np.zeros((num_examples, 1))
        for example in range(num_examples):
            distance = np.zeros((num_centroids, 1))
            for centroid in range(num_centroids):
                distance_diff = data[example, :] - centroids[centroid, :]
                distance[centroid] = np.sum(distance_diff ** 2)
            # np.argmin() 返回数组中最小值的索引
            closest_centroids_ids[example] = np.argmin(distance)
        return closest_centroids_ids

    @staticmethod
    def centroids_compute(data, closest_centroids_ids, num_clusters):
        """ 进行质心更新，即重新计算质心点 """
        num_features = data.shape[1]
        centroids = np.zeros((num_clusters, num_features))
        for cluster in range(num_clusters):
            cluster_data = data[(closest_centroids_ids == cluster).flatten(), :]
            centroids[cluster] = np.mean(cluster_data, axis=0)
        return centroids
